"""
Boussinesq Flow Solver

Implementation of the Boussinesq approximation for natural convection flows.
Handles variable density effects through buoyancy forces while maintaining
incompressibility in the continuity equation.

Key features:
- Natural convection with buoyancy forces
- Temperature-density coupling via thermal expansion
- Integration with heat transfer equations
- Support for both 2D and 3D domains

The Boussinesq approximation:
- Density: ρ = ρ₀[1 - β(T - T₀)]
- Buoyancy force: f_b = -ρg = -ρ₀[1 - β(T - T₀)]g
- Momentum equations include buoyancy source terms
"""

using LinearAlgebra
using SparseArrays

"""
    BoussinesqOptions

Configuration options for Boussinesq natural convection simulations.

# Fields
- `reference_temperature::Float64`: Reference temperature T₀ (K)
- `thermal_expansion::Float64`: Thermal expansion coefficient β (1/K)  
- `gravitational_acceleration::Vector{Float64}`: Gravity vector g (m/s²)
- `density_reference::Float64`: Reference density ρ₀ (kg/m³)
- `prandtl_number::Float64`: Prandtl number Pr = ν/α
"""
Base.@kwdef struct BoussinesqOptions
    reference_temperature::Float64 = 293.15    # K (20°C)
    thermal_expansion::Float64 = 2.1e-4        # 1/K (typical for air)
    gravitational_acceleration::Vector{Float64} = [0.0, -9.81]  # m/s² (2D default)
    density_reference::Float64 = 1.225         # kg/m³ (air at STP)
    prandtl_number::Float64 = 0.71             # Air at room temperature
end

"""
    BoussinesqResult

Results from Boussinesq natural convection solver.

# Fields
- `u::Array`: x-velocity field
- `v::Array`: y-velocity field  
- `w::Union{Array,Nothing}`: z-velocity field (3D only)
- `p::Array`: Pressure field
- `T::Array`: Temperature field
- `x::Vector{Float64}`: x-coordinates
- `y::Vector{Float64}`: y-coordinates
- `z::Union{Vector{Float64},Nothing}`: z-coordinates (3D only)
- `converged::Bool`: Whether solver converged
- `iterations::Int`: Number of time steps performed
- `residual_norm::Float64`: Final residual norm
- `solve_time::Float64`: Total solve time (seconds)
- `convergence_history::Vector{Float64}`: Residual history per iteration
- `buoyancy_force::Array`: Computed buoyancy forces
- `rayleigh_number::Float64`: Computed Rayleigh number
- `nusselt_number::Float64`: Computed Nusselt number (if available)
- `options::NSOptions`: Base solver options used
- `boussinesq_options::BoussinesqOptions`: Boussinesq-specific options
"""
Base.@kwdef struct BoussinesqResult
    u::Array
    v::Array
    w::Union{Array,Nothing}
    p::Array
    T::Array
    x::Vector{Float64}
    y::Vector{Float64}
    z::Union{Vector{Float64},Nothing}
    converged::Bool
    iterations::Int
    residual_norm::Float64
    solve_time::Float64
    convergence_history::Vector{Float64}
    buoyancy_force::Array
    rayleigh_number::Float64
    nusselt_number::Float64 = 0.0
    options::NSOptions
    boussinesq_options::BoussinesqOptions
end

"""
    solve_boussinesq_2d(options::NSOptions, boussinesq::BoussinesqOptions) -> BoussinesqResult

Solve 2D Boussinesq natural convection equations using spectral element methods.

Solves the coupled system:
- Momentum equations with buoyancy forces
- Continuity equation for incompressible flow  
- Energy equation for temperature transport

# Arguments
- `options::NSOptions`: Base Navier-Stokes solver options
- `boussinesq::BoussinesqOptions`: Boussinesq approximation parameters

# Returns
- `BoussinesqResult`: Complete solution including temperature and buoyancy effects
"""
function solve_boussinesq_2d(options::NSOptions, boussinesq::BoussinesqOptions)
    if options.verbose
        println("🌡️  Starting 2D Boussinesq Natural Convection Solver")
        println("Configuration: T₀=$(boussinesq.reference_temperature)K, β=$(boussinesq.thermal_expansion)/K")
        println("Gravity: g=$(boussinesq.gravitational_acceleration) m/s²")
    end
    
    start_time = time()
    
    # Setup domain and grid
    multidomain = create_multidomain(options.n_block, options.N, 0.1, 2)
    n = options.N
    n_block = options.n_block
    
    # Create coordinate arrays 
    x = LinRange(-1, 1, n_block * n + 1)
    y = LinRange(-1, 1, n_block * n + 1)
    nx, ny = length(x), length(y)
    
    # Initialize fields
    u = zeros(nx, ny)
    v = zeros(nx, ny)
    p = zeros(nx, ny)
    T = zeros(nx, ny)
    buoyancy_force = zeros(nx, ny, 2)  # 2D buoyancy force vectors
    
    # Initialize temperature field (differentially heated cavity example)
    # Hot wall on left (x=-1), cold wall on right (x=1)
    T_hot = boussinesq.reference_temperature + 10.0  # Hot wall +10K
    T_cold = boussinesq.reference_temperature - 10.0 # Cold wall -10K
    
    for (i, xi) in enumerate(x), (j, yj) in enumerate(y)
        # Linear temperature profile as initial condition
        T[i,j] = T_cold + (T_hot - T_cold) * (xi + 1) / 2
        
        # Apply boundary conditions
        if i == 1          # Left wall (hot)
            T[i,j] = T_hot
        elseif i == nx     # Right wall (cold)
            T[i,j] = T_cold
        end
    end
    
    # Time stepping parameters
    dt = compute_boussinesq_timestep(options, boussinesq, x, y)
    max_steps = min(options.max_steps, 1000)  # Limit for demonstration
    
    convergence_history = Float64[]
    
    # Main time stepping loop
    for step in 1:max_steps
        # Compute buoyancy forces
        compute_buoyancy_force!(buoyancy_force, T, 
                               boussinesq.density_reference,
                               boussinesq.thermal_expansion,
                               boussinesq.gravitational_acceleration,
                               boussinesq.reference_temperature)
        
        # Solve momentum equations with buoyancy
        solve_momentum_with_buoyancy!(u, v, p, buoyancy_force, dt, options.nu, x, y)
        
        # Solve energy equation
        solve_energy_equation!(T, u, v, dt, boussinesq.prandtl_number * options.nu, x, y)
        
        # Apply boundary conditions
        apply_boussinesq_boundary_conditions!(u, v, T, boussinesq, x, y)
        
        # Check convergence
        residual = compute_boussinesq_residual(u, v, T)
        push!(convergence_history, residual)
        
        if options.verbose && step % 100 == 0
            println("Step $step: Residual = $(residual)")
        end
        
        if residual < options.tol
            if options.verbose
                println("Converged after $step steps")
            end
            break
        end
    end
    
    # Compute final diagnostics
    rayleigh_number = compute_rayleigh_number(boussinesq, maximum(T) - minimum(T), 2.0)
    nusselt_number = compute_nusselt_number_2d(T, x, y)
    
    solve_time = time() - start_time
    final_residual = isempty(convergence_history) ? 0.0 : convergence_history[end]
    
    return BoussinesqResult(
        u = u,
        v = v,
        w = nothing,
        p = p,
        T = T,
        x = collect(x),
        y = collect(y),
        z = nothing,
        converged = final_residual < options.tol,
        iterations = length(convergence_history),
        residual_norm = final_residual,
        solve_time = solve_time,
        convergence_history = convergence_history,
        buoyancy_force = buoyancy_force,
        rayleigh_number = rayleigh_number,
        nusselt_number = nusselt_number,
        options = options,
        boussinesq_options = boussinesq
    )
end

"""
    solve_boussinesq_3d(options::NSOptions, boussinesq::BoussinesqOptions) -> BoussinesqResult

Solve 3D Boussinesq natural convection equations using spectral element methods.

Extends 2D capabilities to full 3D with:
- 3D buoyancy force computations
- 3D energy equation solution
- 3D boundary condition handling

# Arguments
- `options::NSOptions`: Base Navier-Stokes solver options (must have dim=3)
- `boussinesq::BoussinesqOptions`: Boussinesq approximation parameters

# Returns
- `BoussinesqResult`: Complete 3D solution including temperature and buoyancy effects
"""
function solve_boussinesq_3d(options::NSOptions, boussinesq::BoussinesqOptions)
    if options.dim != 3
        throw(ArgumentError("3D Boussinesq solver requires options.dim = 3"))
    end
    
    # Ensure 3D gravity vector
    g_3d = length(boussinesq.gravitational_acceleration) == 3 ? 
           boussinesq.gravitational_acceleration :
           [0.0, 0.0, -9.81]  # Default 3D gravity
    
    boussinesq_3d = BoussinesqOptions(
        reference_temperature = boussinesq.reference_temperature,
        thermal_expansion = boussinesq.thermal_expansion,
        gravitational_acceleration = g_3d,
        density_reference = boussinesq.density_reference,
        prandtl_number = boussinesq.prandtl_number
    )
    
    if options.verbose
        println("🌡️  Starting 3D Boussinesq Natural Convection Solver")
        println("Configuration: T₀=$(boussinesq_3d.reference_temperature)K, β=$(boussinesq_3d.thermal_expansion)/K")
        println("Gravity: g=$(boussinesq_3d.gravitational_acceleration) m/s²")
    end
    
    start_time = time()
    
    # Setup 3D domain and grid
    multidomain = create_multidomain(options.n_block, options.N, 0.1, 3)
    n = options.N
    n_block = options.n_block
    
    # Create 3D coordinate arrays 
    x = LinRange(-1, 1, n_block * n + 1)
    y = LinRange(-1, 1, n_block * n + 1)  
    z = LinRange(-1, 1, n_block * n + 1)
    nx, ny, nz = length(x), length(y), length(z)
    
    # Initialize 3D fields
    u = zeros(nx, ny, nz)
    v = zeros(nx, ny, nz)
    w = zeros(nx, ny, nz)
    p = zeros(nx, ny, nz)
    T = zeros(nx, ny, nz)
    buoyancy_force = zeros(nx, ny, nz, 3)  # 3D buoyancy force vectors
    
    # Initialize 3D temperature field
    T_hot = boussinesq_3d.reference_temperature + 10.0
    T_cold = boussinesq_3d.reference_temperature - 10.0
    
    for (i, xi) in enumerate(x), (j, yj) in enumerate(y), (k, zk) in enumerate(z)
        # Linear temperature profile as initial condition (heated from bottom)
        T[i,j,k] = T_cold + (T_hot - T_cold) * (zk + 1) / 2
        
        # Apply boundary conditions
        if k == 1          # Bottom wall (hot)
            T[i,j,k] = T_hot
        elseif k == nz     # Top wall (cold)
            T[i,j,k] = T_cold
        end
    end
    
    # Time stepping parameters
    dt = compute_boussinesq_timestep_3d(options, boussinesq_3d, x, y, z)
    max_steps = min(options.max_steps, 500)  # Reduced for 3D computational cost
    
    convergence_history = Float64[]
    
    # Main 3D time stepping loop
    for step in 1:max_steps
        # Compute 3D buoyancy forces
        compute_buoyancy_force_3d!(buoyancy_force, T, 
                                  boussinesq_3d.density_reference,
                                  boussinesq_3d.thermal_expansion,
                                  boussinesq_3d.gravitational_acceleration,
                                  boussinesq_3d.reference_temperature)
        
        # Solve 3D momentum equations with buoyancy
        solve_momentum_with_buoyancy_3d!(u, v, w, p, buoyancy_force, dt, options.nu, x, y, z)
        
        # Solve 3D energy equation
        solve_energy_equation_3d!(T, u, v, w, dt, boussinesq_3d.prandtl_number * options.nu, x, y, z)
        
        # Apply 3D boundary conditions
        apply_boussinesq_boundary_conditions_3d!(u, v, w, T, boussinesq_3d, x, y, z)
        
        # Check convergence
        residual = compute_boussinesq_residual_3d(u, v, w, T)
        push!(convergence_history, residual)
        
        if options.verbose && step % 50 == 0
            println("Step $step: Residual = $(residual)")
        end
        
        if residual < options.tol
            if options.verbose
                println("Converged after $step steps")
            end
            break
        end
    end
    
    # Compute final diagnostics
    rayleigh_number = compute_rayleigh_number(boussinesq_3d, maximum(T) - minimum(T), 2.0)
    nusselt_number = compute_nusselt_number_3d(T, x, y, z)
    
    solve_time = time() - start_time
    final_residual = isempty(convergence_history) ? 0.0 : convergence_history[end]
    
    return BoussinesqResult(
        u = u,
        v = v,
        w = w,
        p = p,
        T = T,
        x = collect(x),
        y = collect(y),
        z = collect(z),
        converged = final_residual < options.tol,
        iterations = length(convergence_history),
        residual_norm = final_residual,
        solve_time = solve_time,
        convergence_history = convergence_history,
        buoyancy_force = buoyancy_force,
        rayleigh_number = rayleigh_number,
        nusselt_number = nusselt_number,
        options = options,
        boussinesq_options = boussinesq_3d
    )
end

"""
    compute_buoyancy_force!(f_buoy, T, ρ_ref, β, g, T_ref)

Compute buoyancy forces using the Boussinesq approximation.

Buoyancy force: f_b = -ρg = -ρ₀[1 - β(T - T₀)]g

# Arguments
- `f_buoy::Array`: Output buoyancy force array (modified in-place)
- `T::Array`: Temperature field
- `ρ_ref::Float64`: Reference density ρ₀
- `β::Float64`: Thermal expansion coefficient  
- `g::Vector{Float64}`: Gravitational acceleration vector
- `T_ref::Float64`: Reference temperature T₀
"""
function compute_buoyancy_force!(f_buoy::Array{Float64,3}, T::Array{Float64,2}, 
                                 ρ_ref::Float64, β::Float64, g::Vector{Float64}, T_ref::Float64)
    nx, ny = size(T)
    
    for i in 1:nx, j in 1:ny
        # Density variation: ρ = ρ₀[1 - β(T - T₀)]
        ρ_local = ρ_ref * (1.0 - β * (T[i,j] - T_ref))
        
        # Buoyancy force: f_b = -ρg
        for k in 1:2  # 2D force components
            f_buoy[i,j,k] = -ρ_local * g[k]
        end
    end
end

"""
    compute_buoyancy_force_3d!(f_buoy, T, ρ_ref, β, g, T_ref)

Compute 3D buoyancy forces using the Boussinesq approximation.
"""
function compute_buoyancy_force_3d!(f_buoy::Array{Float64,4}, T::Array{Float64,3}, 
                                   ρ_ref::Float64, β::Float64, g::Vector{Float64}, T_ref::Float64)
    nx, ny, nz = size(T)
    
    for i in 1:nx, j in 1:ny, k in 1:nz
        # Density variation: ρ = ρ₀[1 - β(T - T₀)]
        ρ_local = ρ_ref * (1.0 - β * (T[i,j,k] - T_ref))
        
        # Buoyancy force: f_b = -ρg  
        for d in 1:3  # 3D force components
            f_buoy[i,j,k,d] = -ρ_local * g[d]
        end
    end
end

"""
    solve_momentum_with_buoyancy!(u, v, p, f_buoy, dt, nu, x, y)

Solve momentum equations including buoyancy source terms.

This is a simplified implementation using explicit time stepping.
Production code would use spectral element discretization.
"""
function solve_momentum_with_buoyancy!(u::Array{Float64,2}, v::Array{Float64,2}, 
                                      p::Array{Float64,2}, f_buoy::Array{Float64,3},
                                      dt::Float64, nu::Float64, x::AbstractVector, y::AbstractVector)
    nx, ny = size(u)
    dx = length(x) > 1 ? (x[2] - x[1]) : 1.0
    dy = length(y) > 1 ? (y[2] - y[1]) : 1.0
    
    # Simple finite difference momentum update with buoyancy
    # This is a placeholder - full SEM implementation would be more complex
    u_old = copy(u)
    v_old = copy(v)
    
    for i in 2:(nx-1), j in 2:(ny-1)
        # Simplified momentum equations with buoyancy forces
        # ∂u/∂t = -u∂u/∂x - v∂u/∂y - ∂p/∂x + ν∇²u + f_buoy_x
        # ∂v/∂t = -u∂v/∂x - v∂v/∂y - ∂p/∂y + ν∇²v + f_buoy_y
        
        # Convective terms (simplified)
        dudt_conv = -u_old[i,j] * (u_old[i+1,j] - u_old[i-1,j])/(2*dx) - 
                    v_old[i,j] * (u_old[i,j+1] - u_old[i,j-1])/(2*dy)
                    
        dvdt_conv = -u_old[i,j] * (v_old[i+1,j] - v_old[i-1,j])/(2*dx) - 
                    v_old[i,j] * (v_old[i,j+1] - v_old[i,j-1])/(2*dy)
        
        # Pressure gradients (simplified)
        dpdx = (p[i+1,j] - p[i-1,j])/(2*dx)
        dpdy = (p[i,j+1] - p[i,j-1])/(2*dy)
        
        # Viscous terms (simplified Laplacian)
        d2udx2 = (u_old[i+1,j] - 2*u_old[i,j] + u_old[i-1,j])/(dx^2)
        d2udy2 = (u_old[i,j+1] - 2*u_old[i,j] + u_old[i,j-1])/(dy^2)
        d2vdx2 = (v_old[i+1,j] - 2*v_old[i,j] + v_old[i-1,j])/(dx^2)
        d2vdy2 = (v_old[i,j+1] - 2*v_old[i,j] + v_old[i,j-1])/(dy^2)
        
        # Update with buoyancy forces
        u[i,j] = u_old[i,j] + dt * (dudt_conv - dpdx + nu*(d2udx2 + d2udy2) + f_buoy[i,j,1])
        v[i,j] = v_old[i,j] + dt * (dvdt_conv - dpdy + nu*(d2vdx2 + d2vdy2) + f_buoy[i,j,2])
    end
end

"""
    solve_momentum_with_buoyancy_3d!(u, v, w, p, f_buoy, dt, nu, x, y, z)

Solve 3D momentum equations including buoyancy source terms.
"""
function solve_momentum_with_buoyancy_3d!(u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                                         p::Array{Float64,3}, f_buoy::Array{Float64,4},
                                         dt::Float64, nu::Float64, x::AbstractVector, 
                                         y::AbstractVector, z::AbstractVector)
    nx, ny, nz = size(u)
    dx = length(x) > 1 ? (x[2] - x[1]) : 1.0
    dy = length(y) > 1 ? (y[2] - y[1]) : 1.0
    dz = length(z) > 1 ? (z[2] - z[1]) : 1.0
    
    # Simple 3D finite difference momentum update with buoyancy
    u_old = copy(u)
    v_old = copy(v)
    w_old = copy(w)
    
    for i in 2:(nx-1), j in 2:(ny-1), k in 2:(nz-1)
        # 3D momentum equations with buoyancy forces
        
        # Simplified convective terms
        dudt_conv = -u_old[i,j,k] * (u_old[i+1,j,k] - u_old[i-1,j,k])/(2*dx) - 
                    v_old[i,j,k] * (u_old[i,j+1,k] - u_old[i,j-1,k])/(2*dy) -
                    w_old[i,j,k] * (u_old[i,j,k+1] - u_old[i,j,k-1])/(2*dz)
                    
        dvdt_conv = -u_old[i,j,k] * (v_old[i+1,j,k] - v_old[i-1,j,k])/(2*dx) - 
                    v_old[i,j,k] * (v_old[i,j+1,k] - v_old[i,j-1,k])/(2*dy) -
                    w_old[i,j,k] * (v_old[i,j,k+1] - v_old[i,j,k-1])/(2*dz)
                    
        dwdt_conv = -u_old[i,j,k] * (w_old[i+1,j,k] - w_old[i-1,j,k])/(2*dx) - 
                    v_old[i,j,k] * (w_old[i,j+1,k] - w_old[i,j-1,k])/(2*dy) -
                    w_old[i,j,k] * (w_old[i,j,k+1] - w_old[i,j,k-1])/(2*dz)
        
        # Pressure gradients
        dpdx = (p[i+1,j,k] - p[i-1,j,k])/(2*dx)
        dpdy = (p[i,j+1,k] - p[i,j-1,k])/(2*dy)
        dpdz = (p[i,j,k+1] - p[i,j,k-1])/(2*dz)
        
        # 3D viscous terms
        laplacian_u = (u_old[i+1,j,k] - 2*u_old[i,j,k] + u_old[i-1,j,k])/(dx^2) +
                      (u_old[i,j+1,k] - 2*u_old[i,j,k] + u_old[i,j-1,k])/(dy^2) +
                      (u_old[i,j,k+1] - 2*u_old[i,j,k] + u_old[i,j,k-1])/(dz^2)
                      
        laplacian_v = (v_old[i+1,j,k] - 2*v_old[i,j,k] + v_old[i-1,j,k])/(dx^2) +
                      (v_old[i,j+1,k] - 2*v_old[i,j,k] + v_old[i,j-1,k])/(dy^2) +
                      (v_old[i,j,k+1] - 2*v_old[i,j,k] + v_old[i,j,k-1])/(dz^2)
                      
        laplacian_w = (w_old[i+1,j,k] - 2*w_old[i,j,k] + w_old[i-1,j,k])/(dx^2) +
                      (w_old[i,j+1,k] - 2*w_old[i,j,k] + w_old[i,j-1,k])/(dy^2) +
                      (w_old[i,j,k+1] - 2*w_old[i,j,k] + w_old[i,j,k-1])/(dz^2)
        
        # Update with buoyancy forces
        u[i,j,k] = u_old[i,j,k] + dt * (dudt_conv - dpdx + nu*laplacian_u + f_buoy[i,j,k,1])
        v[i,j,k] = v_old[i,j,k] + dt * (dvdt_conv - dpdy + nu*laplacian_v + f_buoy[i,j,k,2])
        w[i,j,k] = w_old[i,j,k] + dt * (dwdt_conv - dpdz + nu*laplacian_w + f_buoy[i,j,k,3])
    end
end

"""
    solve_energy_equation!(T, u, v, dt, alpha, x, y)

Solve the energy equation: ∂T/∂t + u·∇T = α∇²T

Where α is the thermal diffusivity.
"""
function solve_energy_equation!(T::Array{Float64,2}, u::Array{Float64,2}, v::Array{Float64,2},
                                dt::Float64, alpha::Float64, x::AbstractVector, y::AbstractVector)
    nx, ny = size(T)
    dx = length(x) > 1 ? (x[2] - x[1]) : 1.0
    dy = length(y) > 1 ? (y[2] - y[1]) : 1.0
    
    T_old = copy(T)
    
    for i in 2:(nx-1), j in 2:(ny-1)
        # Convective terms: u·∇T
        dTdx = (T_old[i+1,j] - T_old[i-1,j])/(2*dx)
        dTdy = (T_old[i,j+1] - T_old[i,j-1])/(2*dy)
        convection = u[i,j] * dTdx + v[i,j] * dTdy
        
        # Diffusive terms: α∇²T
        d2Tdx2 = (T_old[i+1,j] - 2*T_old[i,j] + T_old[i-1,j])/(dx^2)
        d2Tdy2 = (T_old[i,j+1] - 2*T_old[i,j] + T_old[i,j-1])/(dy^2)
        diffusion = alpha * (d2Tdx2 + d2Tdy2)
        
        # Update temperature
        T[i,j] = T_old[i,j] + dt * (-convection + diffusion)
    end
end

"""
    solve_energy_equation_3d!(T, u, v, w, dt, alpha, x, y, z)

Solve the 3D energy equation.
"""
function solve_energy_equation_3d!(T::Array{Float64,3}, u::Array{Float64,3}, 
                                  v::Array{Float64,3}, w::Array{Float64,3},
                                  dt::Float64, alpha::Float64, x::AbstractVector, 
                                  y::AbstractVector, z::AbstractVector)
    nx, ny, nz = size(T)
    dx = length(x) > 1 ? (x[2] - x[1]) : 1.0
    dy = length(y) > 1 ? (y[2] - y[1]) : 1.0
    dz = length(z) > 1 ? (z[2] - z[1]) : 1.0
    
    T_old = copy(T)
    
    for i in 2:(nx-1), j in 2:(ny-1), k in 2:(nz-1)
        # 3D convective terms
        dTdx = (T_old[i+1,j,k] - T_old[i-1,j,k])/(2*dx)
        dTdy = (T_old[i,j+1,k] - T_old[i,j-1,k])/(2*dy)
        dTdz = (T_old[i,j,k+1] - T_old[i,j,k-1])/(2*dz)
        convection = u[i,j,k] * dTdx + v[i,j,k] * dTdy + w[i,j,k] * dTdz
        
        # 3D diffusive terms
        laplacian_T = (T_old[i+1,j,k] - 2*T_old[i,j,k] + T_old[i-1,j,k])/(dx^2) +
                      (T_old[i,j+1,k] - 2*T_old[i,j,k] + T_old[i,j-1,k])/(dy^2) +
                      (T_old[i,j,k+1] - 2*T_old[i,j,k] + T_old[i,j,k-1])/(dz^2)
        diffusion = alpha * laplacian_T
        
        # Update temperature
        T[i,j,k] = T_old[i,j,k] + dt * (-convection + diffusion)
    end
end

"""
    apply_boussinesq_boundary_conditions!(u, v, T, boussinesq, x, y)

Apply boundary conditions for Boussinesq natural convection problems.
"""
function apply_boussinesq_boundary_conditions!(u::Array{Float64,2}, v::Array{Float64,2}, 
                                              T::Array{Float64,2}, boussinesq::BoussinesqOptions,
                                              x::AbstractVector, y::AbstractVector)
    nx, ny = size(u)
    
    # No-slip conditions on all walls
    u[1,:] .= 0.0    # Left wall
    u[nx,:] .= 0.0   # Right wall
    u[:,1] .= 0.0    # Bottom wall
    u[:,ny] .= 0.0   # Top wall
    
    v[1,:] .= 0.0    # Left wall
    v[nx,:] .= 0.0   # Right wall
    v[:,1] .= 0.0    # Bottom wall
    v[:,ny] .= 0.0   # Top wall
    
    # Temperature boundary conditions (differentially heated cavity)
    T_hot = boussinesq.reference_temperature + 10.0
    T_cold = boussinesq.reference_temperature - 10.0
    
    T[1,:] .= T_hot     # Left wall (hot)
    T[nx,:] .= T_cold   # Right wall (cold)
    # Insulated top and bottom walls (zero flux) - handled by natural evolution
end

"""
    apply_boussinesq_boundary_conditions_3d!(u, v, w, T, boussinesq, x, y, z)

Apply 3D boundary conditions for Boussinesq natural convection.
"""
function apply_boussinesq_boundary_conditions_3d!(u::Array{Float64,3}, v::Array{Float64,3}, 
                                                 w::Array{Float64,3}, T::Array{Float64,3}, 
                                                 boussinesq::BoussinesqOptions,
                                                 x::AbstractVector, y::AbstractVector, z::AbstractVector)
    nx, ny, nz = size(u)
    
    # No-slip conditions on all walls
    u[1,:,:] .= 0.0    # Left wall (x)
    u[nx,:,:] .= 0.0   # Right wall (x)
    u[:,1,:] .= 0.0    # Front wall (y)
    u[:,ny,:] .= 0.0   # Back wall (y)
    u[:,:,1] .= 0.0    # Bottom wall (z)
    u[:,:,nz] .= 0.0   # Top wall (z)
    
    v[1,:,:] .= 0.0    # Left wall (x)
    v[nx,:,:] .= 0.0   # Right wall (x)
    v[:,1,:] .= 0.0    # Front wall (y)
    v[:,ny,:] .= 0.0   # Back wall (y)
    v[:,:,1] .= 0.0    # Bottom wall (z)
    v[:,:,nz] .= 0.0   # Top wall (z)
    
    w[1,:,:] .= 0.0    # Left wall (x)
    w[nx,:,:] .= 0.0   # Right wall (x)
    w[:,1,:] .= 0.0    # Front wall (y)
    w[:,ny,:] .= 0.0   # Back wall (y)
    w[:,:,1] .= 0.0    # Bottom wall (z)
    w[:,:,nz] .= 0.0   # Top wall (z)
    
    # Temperature boundary conditions (Rayleigh-Bénard convection)
    T_hot = boussinesq.reference_temperature + 10.0
    T_cold = boussinesq.reference_temperature - 10.0
    
    T[:,:,1] .= T_hot     # Bottom wall (hot)
    T[:,:,nz] .= T_cold   # Top wall (cold)
    # Insulated side walls - handled by natural evolution
end

# Utility functions

"""
    compute_boussinesq_timestep(options, boussinesq, x, y) -> Float64

Compute stable time step for Boussinesq simulation.
"""
function compute_boussinesq_timestep(options::NSOptions, boussinesq::BoussinesqOptions, 
                                    x::AbstractVector, y::AbstractVector)
    dx = length(x) > 1 ? abs(x[2] - x[1]) : 1.0
    dy = length(y) > 1 ? abs(y[2] - y[1]) : 1.0
    h_min = min(dx, dy)
    
    # CFL condition considering both momentum and energy equations
    # dt ≤ CFL * h²/(2*max(ν, α)) for stability
    alpha = boussinesq.prandtl_number * options.nu  # Thermal diffusivity
    max_diffusivity = max(options.nu, alpha)
    
    dt_diffusion = options.cfl * h_min^2 / (2 * max_diffusivity)
    dt_convection = options.cfl * h_min / 1.0  # Assuming max velocity ~ 1
    
    return min(dt_diffusion, dt_convection)
end

"""
    compute_boussinesq_timestep_3d(options, boussinesq, x, y, z) -> Float64

Compute stable time step for 3D Boussinesq simulation.
"""
function compute_boussinesq_timestep_3d(options::NSOptions, boussinesq::BoussinesqOptions, 
                                       x::AbstractVector, y::AbstractVector, z::AbstractVector)
    dx = length(x) > 1 ? abs(x[2] - x[1]) : 1.0
    dy = length(y) > 1 ? abs(y[2] - y[1]) : 1.0
    dz = length(z) > 1 ? abs(z[2] - z[1]) : 1.0
    h_min = min(dx, dy, dz)
    
    alpha = boussinesq.prandtl_number * options.nu
    max_diffusivity = max(options.nu, alpha)
    
    # 3D stability condition is more restrictive
    dt_diffusion = options.cfl * h_min^2 / (6 * max_diffusivity)
    dt_convection = options.cfl * h_min / 1.0
    
    return min(dt_diffusion, dt_convection)
end

"""
    compute_boussinesq_residual(u, v, T) -> Float64

Compute residual norm for convergence checking.
"""
function compute_boussinesq_residual(u::Array{Float64,2}, v::Array{Float64,2}, T::Array{Float64,2})
    return sqrt(sum(u.^2) + sum(v.^2) + sum(T.^2)) / (length(u) + length(v) + length(T))
end

"""
    compute_boussinesq_residual_3d(u, v, w, T) -> Float64

Compute 3D residual norm for convergence checking.
"""
function compute_boussinesq_residual_3d(u::Array{Float64,3}, v::Array{Float64,3}, 
                                       w::Array{Float64,3}, T::Array{Float64,3})
    return sqrt(sum(u.^2) + sum(v.^2) + sum(w.^2) + sum(T.^2)) / (length(u) + length(v) + length(w) + length(T))
end

"""
    compute_rayleigh_number(boussinesq, ΔT, L) -> Float64

Compute the Rayleigh number: Ra = gβΔTL³/(να)
"""
function compute_rayleigh_number(boussinesq::BoussinesqOptions, ΔT::Float64, L::Float64)
    g_magnitude = norm(boussinesq.gravitational_acceleration)
    α = 1.0 / boussinesq.prandtl_number  # Thermal diffusivity ratio
    
    # Ra = gβΔTL³/(να) where ν is kinematic viscosity, α is thermal diffusivity
    # For this implementation, we assume ν and α are given through Prandtl number
    return g_magnitude * boussinesq.thermal_expansion * ΔT * L^3 / α
end

"""
    compute_nusselt_number_2d(T, x, y) -> Float64

Compute average Nusselt number for 2D natural convection.
"""
function compute_nusselt_number_2d(T::Array{Float64,2}, x::AbstractVector, y::AbstractVector)
    nx, ny = size(T)
    
    if nx < 3 || ny < 3
        return 1.0  # Conduction limit
    end
    
    # Simplified Nusselt number calculation
    # Nu = hL/k where h is heat transfer coefficient
    # For natural convection: Nu ≈ average temperature gradient at walls
    
    dx = abs(x[2] - x[1])
    
    # Average heat flux at hot wall (left wall, i=1)
    heat_flux_hot = 0.0
    for j in 2:(ny-1)
        dTdx_wall = (T[2,j] - T[1,j]) / dx  # Temperature gradient at wall
        heat_flux_hot += abs(dTdx_wall)
    end
    heat_flux_hot /= (ny - 2)
    
    # Reference: conductive heat flux for temperature difference
    ΔT = maximum(T) - minimum(T)
    L = abs(x[end] - x[1])  # Characteristic length
    conductive_flux = ΔT / L
    
    # Nusselt number
    Nu = heat_flux_hot / conductive_flux
    
    return max(Nu, 1.0)  # Nu ≥ 1 for natural convection
end

"""
    compute_nusselt_number_3d(T, x, y, z) -> Float64

Compute average Nusselt number for 3D natural convection.
"""
function compute_nusselt_number_3d(T::Array{Float64,3}, x::AbstractVector, 
                                  y::AbstractVector, z::AbstractVector)
    nx, ny, nz = size(T)
    
    if nx < 3 || ny < 3 || nz < 3
        return 1.0  # Conduction limit
    end
    
    # For Rayleigh-Bénard convection, compute Nu at bottom heated wall (k=1)
    dz = abs(z[2] - z[1])
    
    # Average heat flux at heated bottom wall
    heat_flux_bottom = 0.0
    count = 0
    
    for i in 2:(nx-1), j in 2:(ny-1)
        dTdz_wall = (T[i,j,2] - T[i,j,1]) / dz
        heat_flux_bottom += abs(dTdz_wall)
        count += 1
    end
    
    if count > 0
        heat_flux_bottom /= count
    end
    
    # Reference conductive flux
    ΔT = maximum(T) - minimum(T)
    L = abs(z[end] - z[1])
    conductive_flux = ΔT / L
    
    # Nusselt number
    Nu = heat_flux_bottom / conductive_flux
    
    return max(Nu, 1.0)
end